# Ultralytics YOLO 🚀, AGPL-3.0 license
"""Advanced Loss Functions for YOLOv8 Custom Training."""

import torch
import torch.nn as nn
import torch.nn.functional as F


def bbox_wiou(pred_boxes, gt_boxes, xywh=True, eps=1e-7):
    """
    计算 Wise-IoU (WIoU) Loss
    
    WIoU通过动态调整质量权重来改进IoU Loss:
    - 高质量预测框(IoU>0.7): 较小梯度，避免过拟合
    - 低质量预测框(IoU<0.5): 较大梯度，加速收敛
    
    Args:
        pred_boxes (torch.Tensor): 预测框 (N, 4)
        gt_boxes (torch.Tensor): 真实框 (N, 4)
        xywh (bool): 是否为xywh格式
        eps (float): 防止除零的小值
    
    Returns:
        torch.Tensor: WIoU损失值
    """
    # 转换为 xyxy 格式
    if xywh:
        (x1, y1, w1, h1), (x2, y2, w2, h2) = pred_boxes.chunk(4, -1), gt_boxes.chunk(4, -1)
        w1_, h1_, w2_, h2_ = w1 / 2, h1 / 2, w2 / 2, h2 / 2
        b1_x1, b1_x2, b1_y1, b1_y2 = x1 - w1_, x1 + w1_, y1 - h1_, y1 + h1_
        b2_x1, b2_x2, b2_y1, b2_y2 = x2 - w2_, x2 + w2_, y2 - h2_, y2 + h2_
    else:
        b1_x1, b1_y1, b1_x2, b1_y2 = pred_boxes.chunk(4, -1)
        b2_x1, b2_y1, b2_x2, b2_y2 = gt_boxes.chunk(4, -1)
        w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
        w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps

    # 计算交集面积
    inter = (b1_x2.minimum(b2_x2) - b1_x1.maximum(b2_x1)).clamp_(0) * \
            (b1_y2.minimum(b2_y2) - b1_y1.maximum(b2_y1)).clamp_(0)

    # 计算并集面积
    union = w1 * h1 + w2 * h2 - inter + eps

    # 基础IoU
    iou = inter / union

    # 计算动态β参数
    # β = (IoU_mean - IoU) / (IoU_max - IoU_mean + ε)
    iou_mean = iou.mean().detach()
    iou_max = iou.max().detach()
    beta = (iou_mean - iou) / (iou_max - iou_mean + eps)
    beta = beta.clamp(-3.0, 3.0)  # 限制范围，避免极端值

    # 计算质量权重
    # quality_weight = exp(-β * (1 - IoU)²)
    quality_weight = torch.exp(-beta * (1.0 - iou).pow(2))

    # WIoU Loss = (1 - IoU) * quality_weight
    loss_wiou = (1.0 - iou) * quality_weight

    return loss_wiou


class WIoULoss(nn.Module):
    """
    Wise-IoU Loss 模块
    
    特点:
    1. 自适应调整不同质量预测框的损失权重
    2. 高质量框: 小梯度，避免过拟合
    3. 低质量框: 大梯度，快速收敛
    """
    
    def __init__(self):
        super().__init__()
    
    def forward(self, pred_boxes, gt_boxes, xywh=True):
        """
        Args:
            pred_boxes: 预测边界框 (N, 4)
            gt_boxes: 真实边界框 (N, 4)
            xywh: 是否为xywh格式
        
        Returns:
            loss: WIoU损失值
        """
        return bbox_wiou(pred_boxes, gt_boxes, xywh=xywh).mean()


class ClassBalancedFocalLoss(nn.Module):
    """
    类别平衡Focal Loss
    
    改进点:
    1. 动态α平衡因子: 根据类别样本数自动调整
    2. 自适应γ聚焦参数: 根据样本难度动态调整
    3. 解决类别不平衡问题，特别适合单类别检测
    
    公式:
    CBFocalLoss = -α_balanced * (1 - p_t)^γ_adaptive * log(p_t)
    
    其中:
    - α_balanced = clip(n_total / (n_class * n_positive), 0.1, 10.0)
    - γ_adaptive = γ_base + k * difficulty_score
    """
    
    def __init__(self, nc=1, gamma=2.0, alpha=0.25, adaptive_gamma=True):
        """
        Args:
            nc (int): 类别数量
            gamma (float): 基础聚焦参数
            alpha (float): 基础平衡因子
            adaptive_gamma (bool): 是否使用自适应gamma
        """
        super().__init__()
        self.nc = nc
        self.gamma = gamma
        self.alpha = alpha
        self.adaptive_gamma = adaptive_gamma
    
    def forward(self, pred_scores, target_scores, reduction='mean'):
        """
        Args:
            pred_scores: 预测分数 (B, N, C) - logits形式
            target_scores: 目标分数 (B, N, C) - 0/1或软标签
            reduction: 'mean', 'sum', 'none'
        
        Returns:
            loss: Focal Loss值
        """
        # 计算预测概率
        pred_prob = pred_scores.sigmoid()
        
        # 计算 p_t: 正样本用预测概率，负样本用1-预测概率
        p_t = target_scores * pred_prob + (1 - target_scores) * (1 - pred_prob)
        
        # 1. 计算类别平衡因子
        n_pos = (target_scores == 1).sum().float()
        n_total = target_scores.numel()
        
        if n_pos > 0:
            # α_balanced = n_total / (n_class * n_positive)
            alpha_balanced = n_total / (self.nc * n_pos)
            alpha_balanced = torch.clamp(alpha_balanced, 0.1, 10.0)  # 限制范围
        else:
            alpha_balanced = torch.tensor(self.alpha, device=pred_scores.device)
        
        # 为正负样本分配不同的α
        alpha_weight = target_scores * alpha_balanced + (1 - target_scores) * (1 - alpha_balanced)
        
        # 2. 计算难度（置信度低的样本更难）
        difficulty = (1 - p_t).clamp(eps=1e-7)  # 避免log(0)
        
        # 3. 自适应gamma
        if self.adaptive_gamma:
            # γ_adaptive = γ_base + k * difficulty_mean
            difficulty_score = difficulty.mean().detach()
            gamma_adaptive = self.gamma + 0.5 * difficulty_score
            gamma_adaptive = torch.clamp(gamma_adaptive, 1.0, 4.0)
        else:
            gamma_adaptive = self.gamma
        
        # 4. 计算Focal权重: (1 - p_t)^γ
        focal_weight = difficulty.pow(gamma_adaptive)
        
        # 5. 计算BCE Loss
        ce_loss = F.binary_cross_entropy_with_logits(
            pred_scores, 
            target_scores, 
            reduction='none'
        )
        
        # 6. 组合: α * focal_weight * ce_loss
        loss = alpha_weight * focal_weight * ce_loss
        
        # 7. Reduction
        if reduction == 'mean':
            return loss.mean()
        elif reduction == 'sum':
            return loss.sum()
        else:
            return loss


class CombinedLoss(nn.Module):
    """
    组合损失: WIoU + 类别平衡Focal Loss + DFL
    
    用于替换YOLOv8的默认损失函数
    """
    
    def __init__(self, nc=1, reg_max=16, use_dfl=True):
        """
        Args:
            nc: 类别数
            reg_max: DFL的最大值
            use_dfl: 是否使用DFL
        """
        super().__init__()
        self.nc = nc
        self.reg_max = reg_max
        self.use_dfl = use_dfl
        
        # WIoU Loss
        self.wiou_loss = WIoULoss()
        
        # 类别平衡Focal Loss
        self.focal_loss = ClassBalancedFocalLoss(nc=nc, gamma=2.0, alpha=0.25)
    
    def forward(self, pred_dist, pred_bboxes, anchor_points, target_bboxes, 
                target_scores, target_scores_sum, fg_mask):
        """
        计算组合损失
        
        Args:
            pred_dist: 预测分布 (B, N, 4*reg_max)
            pred_bboxes: 解码后的预测框 (B, N, 4)
            anchor_points: 锚点 (N, 2)
            target_bboxes: 目标框 (B, N, 4)
            target_scores: 目标分数 (B, N, C)
            target_scores_sum: 目标分数总和
            fg_mask: 前景mask (B, N)
        
        Returns:
            loss_iou: WIoU损失
            loss_dfl: DFL损失
        """
        # 1. WIoU Loss (仅对前景)
        if fg_mask.sum() > 0:
            weight = target_scores.sum(-1)[fg_mask].unsqueeze(-1)
            
            # 计算WIoU
            wiou = bbox_wiou(pred_bboxes[fg_mask], target_bboxes[fg_mask], xywh=False)
            loss_iou = (wiou * weight).sum() / target_scores_sum
        else:
            loss_iou = torch.tensor(0.0, device=pred_bboxes.device)
        
        # 2. DFL Loss
        if self.use_dfl and fg_mask.sum() > 0:
            # 将target_bboxes转换为dist格式
            from ultralytics.utils.tal import bbox2dist
            target_ltrb = bbox2dist(anchor_points, target_bboxes, self.reg_max)
            loss_dfl = self._df_loss(
                pred_dist[fg_mask].view(-1, self.reg_max + 1), 
                target_ltrb[fg_mask]
            ) * weight
            loss_dfl = loss_dfl.sum() / target_scores_sum
        else:
            loss_dfl = torch.tensor(0.0, device=pred_bboxes.device)
        
        return loss_iou, loss_dfl
    
    @staticmethod
    def _df_loss(pred_dist, target):
        """Distribution Focal Loss"""
        tl = target.long()
        tr = tl + 1
        wl = tr - target
        wr = 1 - wl
        return (
            F.cross_entropy(pred_dist, tl.view(-1), reduction="none").view(tl.shape) * wl
            + F.cross_entropy(pred_dist, tr.view(-1), reduction="none").view(tl.shape) * wr
        ).mean(-1, keepdim=True)


# 测试代码
if __name__ == "__main__":
    print("=" * 60)
    print("测试 WIoU Loss")
    print("=" * 60)
    
    # 测试WIoU
    pred = torch.tensor([[10, 10, 20, 20], [15, 15, 25, 25], [5, 5, 15, 15]], dtype=torch.float32)
    gt = torch.tensor([[10, 10, 20, 20], [10, 10, 20, 20], [10, 10, 20, 20]], dtype=torch.float32)
    
    wiou_loss = WIoULoss()
    loss = wiou_loss(pred, gt, xywh=False)
    print(f"WIoU Loss: {loss.item():.4f}")
    
    print("\n" + "=" * 60)
    print("测试 类别平衡Focal Loss")
    print("=" * 60)
    
    # 测试Focal Loss
    pred_scores = torch.randn(2, 100, 1)  # (B, N, C)
    target_scores = torch.zeros(2, 100, 1)
    target_scores[:, :10, :] = 1  # 只有10%是正样本
    
    focal_loss = ClassBalancedFocalLoss(nc=1)
    loss = focal_loss(pred_scores, target_scores)
    print(f"Class Balanced Focal Loss: {loss.item():.4f}")
    
    print("\n✅ 损失函数测试完成！")

